

#libraries and modules
import numpy as np
import matplotlib.pyplot as plt




def f_S(CN):
    #INPUT:
    #CN,            SCS run-off curve number

    #DESCRIPTION:
    #compute potential maximum retention, S

    #OUTPUT:
    #S         potential maximum retention in inches 

    
    S=(1000*CN**-1)-10    #TR55, eq. 2-4 
    
    return S




def f_Ia(S):
    #INPUT:
    #S              potential maximum retention, in inches           

    #DESCRIPTION:
    #compute initial abstraction, Ia

    #OUTPUT:
    #Ia             initial abstraction in inches

    Ia=0.2*S     #TR55, eq. 2-2; empirical relation from many small agricultural watersheds in USA
    

    return Ia


def f_Q(P,CN):
    #INPUT:
    #P              precipitation, in inches
    #CN             SCS run-off curve number
             

    #DESCRIPTION:
    #compute Run-off, Q

    #OUTPUT:
    #Q              run-off in inches


    #calclation of potential maximum retention, S
    S=f_S(CN)

    #calculation of initial abstraction, Ia
    Ia=f_Ia(S)
    
    #calculation of Q

    try:
        #if P is array
        P=np.asarray(P)
        Q=np.where(P>Ia,
                   ((P-Ia)**2) / ((P-Ia)+S),
                   np.zeros(len(P)))                #TR55, eq. 2-1; modified for P<=Ia
    except:
        #if P is not array
        if P>Ia:Q=((P-Ia)**2) / ((P-Ia)+S)
        else:Q=0

    return Q

def plot_Q(P,CN, metric):
    #INPUT:
    #P              precipitation, in inches
    #CN             SCS run-off curve number
    #metric         boolean TRUE=metric, FALSE=english

    #DESCRIPTION
    #plots Q as a function of P

    Q=f_Q(P,CN)
    
    if metric: c=25.4
    else:c=1

    plt.plot(P*c, Q*c, '.-')





    


